"""
Joint mapping utilities for converting between Isaac Sim and Isaac Gym conventions.
"""

import numpy as np


class JointMapping:
    """关节映射类，处理Sim和Gym环境之间的关节顺序转换"""

    def __init__(self):
        # G1机器人的关节映射关系 - 完全复制源代码
        # Sim环境关节顺序（27个关节）- 从源代码复制
        self.sim_joint_names = [
            'left_hip_pitch_joint',
            'right_hip_pitch_joint',
            'waist_yaw_joint',
            'left_hip_roll_joint',
            'right_hip_roll_joint',
            'left_hip_yaw_joint',
            'right_hip_yaw_joint',
            'left_knee_joint',
            'right_knee_joint',
            'left_shoulder_pitch_joint',
            'right_shoulder_pitch_joint',
            'left_ankle_pitch_joint',
            'right_ankle_pitch_joint',
            'left_shoulder_roll_joint',
            'right_shoulder_roll_joint',
            'left_ankle_roll_joint',
            'right_ankle_roll_joint',
            'left_shoulder_yaw_joint',
            'right_shoulder_yaw_joint',
            'left_elbow_joint',
            'right_elbow_joint',
            'left_wrist_roll_joint',
            'right_wrist_roll_joint',
            'left_wrist_pitch_joint',
            'right_wrist_pitch_joint',
            'left_wrist_yaw_joint',
            'right_wrist_yaw_joint',
        ]

        # Gym环境关节顺序（27个关节）- 从源代码复制
        self.gym_joint_names = [
            'left_hip_pitch_joint',
            'left_hip_roll_joint',
            'left_hip_yaw_joint',
            'left_knee_joint',
            'left_ankle_pitch_joint',
            'left_ankle_roll_joint',
            'right_hip_pitch_joint',
            'right_hip_roll_joint',
            'right_hip_yaw_joint',
            'right_knee_joint',
            'right_ankle_pitch_joint',
            'right_ankle_roll_joint',
            'waist_yaw_joint',
            'left_shoulder_pitch_joint',
            'left_shoulder_roll_joint',
            'left_shoulder_yaw_joint',
            'left_elbow_joint',
            'left_wrist_roll_joint',
            'left_wrist_pitch_joint',
            'left_wrist_yaw_joint',
            'right_shoulder_pitch_joint',
            'right_shoulder_roll_joint',
            'right_shoulder_yaw_joint',
            'right_elbow_joint',
            'right_wrist_roll_joint',
            'right_wrist_pitch_joint',
            'right_wrist_yaw_joint',
        ]

        # 创建映射索引
        self.sim2gym_indices = [self.sim_joint_names.index(name) for name in self.gym_joint_names]
        self.gym2sim_indices = [self.gym_joint_names.index(name) for name in self.sim_joint_names]

    def sim2gym(self, sim_positions):
        """从Sim顺序转换到Gym顺序"""
        if len(sim_positions) != 27:
            raise ValueError(f"Sim关节位置数量应为27，实际为{len(sim_positions)}")
        return np.array(sim_positions)[self.sim2gym_indices]

    def gym2sim(self, gym_positions):
        """从Gym顺序转换到Sim顺序"""
        if len(gym_positions) != 27:
            raise ValueError(f"Gym关节位置数量应为27，实际为{len(gym_positions)}")
        return np.array(gym_positions)[self.gym2sim_indices]